execfile('imports.py')

# input directory with pdb fils
in_single = './input/single23/'
in_test = './input/test/'
input_dir = in_single
# output directory will contain pdb files nad pickle files.
out_single = './output/single23/'
out_test = './output/test/'
output_dir = out_single
# output dir for colored pdb files.
bq_single = './bq/single23/'
bq_test = './bq/test/'
bq_dir = bq_single


# what to do:
sandbox = True
test = True
nono = False
proc = True
check = True
write = True
run = True
derive = True
transfer = True
dump = True
color = True

if test:
    input_dir = in_test
    output_dir = out_test
    bq_dir = bq_test

# processing functions
# pre- processing of structures
proc_funcs = ['nohydrogen',
              'nohetero',
              'nomodres',
              'noical']

check_funcs = ['check_lenght',
               'check_bbone',
               'check_space',
               'check_unk',
               'check_lenght']

run_funcs = ['run_areaimol',
             'run_areaimol_asu',
             'run_stride',
             'run_ncont',
             'run_act',
             'run_msms']



def run_ncont(structure, file_):
    model = structure[(0,)]
    out1 = NCONTAtomContacts(file_, model, maxdist =5, xtra_key ='CNT_NCONT_X5')
    out2 = NCONTAtomContacts(file_, model, maxdist =7, xtra_key ='CNT_NCONT_X7')
    return [bool(out1), bool(out2)]

def run_act(structure, file_):
    model = structure[(0,)]
    out1 = ACTAtomContacts(file_, model, contact ='inter 0.0 5.0', xtra_key ='CNT_ACT_X5')
    out2 = ACTAtomContacts(file_, model, contact ='inter 0.0 7.0', xtra_key ='CNT_ACT_X7')
    return [bool(out1), bool(out2)]

def run_areaimol(structure, file_):
    model = structure[(0,)]
    out1 = AreaimolAtomSurface(file_, model)
    out2 = ContactAtomSurface(file_, model)
    return [bool(out1), bool(out2)]

def run_areaimol_asu(structure, file_):
    model = structure[(0,)]
    out1 = ASUAreaimolAtomSurface(model)
    out2 = asu_contacts(model)
    return [bool(out1), bool(out2)]

def run_stride(structure, file_):
    model = structure[(0,)]
    out1 = StrideResidueSurface(file_, model)
    out2 = SecondaryStructure(file_, model)
    return [bool(out1), bool(out2)]

def run_msms(structure, file_):
    model = structure[(0,)]
    out1 = AtomDepth(model)
    out2 = BoolAtomDepth(model)
    return [bool(out1), bool(out2)]

def derive_externals(structure):
    model = structure[(0,)]
    atoms = einput(model, 'A')
    residues = einput(model, 'R')
    chains = einput(model, 'C')
    # preparation
    residues.dispatch('set_coords')

    #AREAIMOL
    model._data_propagate('sum','A','EXP_AREAIMOL_ASA', xtra =True)
    model._data_propagate('sum','A','CNT_AREAIMOL_ASA', xtra =True)
    RelativeSurface(model, 'EXP_AREAIMOL_ASA')
    RelativeSurface(model, 'CNT_AREAIMOL_ASA')
    model._data_propagate('sum','A','ASU_EXP_AREAIMOL_ASA', xtra =True)
    model._data_propagate('sum','A','ASU_CNT_AREAIMOL_ASA', xtra =True)
    RelativeSurface(model, 'ASU_EXP_AREAIMOL_ASA')
    RelativeSurface(model, 'ASU_CNT_AREAIMOL_ASA')

    #ACT
    model._data_propagate('merge','A','CNT_ACT_X5', xtra =True)
    model._data_propagate('merge','A','CNT_ACT_X7', xtra =True)
    find_coordinations(model, 'A', cnt_key ='CNT_ACT_X5', xtra_key = 'CN_ACT_X5')
    find_coordinations(model, 'A', cnt_key ='CNT_ACT_X7', xtra_key = 'CN_ACT_X7')
    model._data_propagate('unique','A','CN_ACT_X5', xtra =True)
    model._data_propagate('unique','A','CN_ACT_X7', xtra =True)

    #NCONT
    model._data_propagate('merge','A','CNT_NCONT_X5', xtra =True)
    model._data_propagate('merge','A','CNT_NCONT_X7', xtra =True)
    find_coordinations(model, 'A', cnt_key ='CNT_NCONT_X5', xtra_key = 'CN_NCONT_X5')
    find_coordinations(model, 'A', cnt_key ='CNT_NCONT_X7', xtra_key = 'CN_NCONT_X7')
    model._data_propagate('unique','A','CN_NCONT_X5', xtra =True)
    model._data_propagate('unique','A','CN_NCONT_X7', xtra =True)

    #Tables..
    ResidueCharge(model)
    ResidueSolvatation(model)
    ResidueGravy(model)
    ResidueEntropy(model)
    ResiduePolarity(model)
    ResidueAA_2(model)
    ResidueAA_5(model)
    ResidueAA_4MURPHY(model)
    ResidueAA_8MURPHY(model)
    ResidueAA_10MURPHY(model)
    ResidueAA_15MURPHY(model)

def transfer_coordination(structure):
    model = structure[(0,)]
    residues = einput(model, 'R')
    asar = residues._select_children(0.0, 'lt', 'CNT_AREAIMOL_ASA', xtra =True).values()
    cntr = residues._select_children({}, 'lt', 'CNT_ACT_X5', xtra =True).values()
    transfer_nearest_xtra(asar, cntr, 'R', 'CN_ACT_X5')
    AttrResidueClusters(asar, xtra_key ='CNT_CLUST', attr ='CN_ACT_X5', xtra =True)


def color_bq(structure, output_dir):
    BeQu("%s%s_%s.pdb" % (output_dir, structure.name, 'cnt'),
         structure, 'R', b_key ='CNT_AREAIMOL_ASA', q_key ='CNT_ACT_X5', q_mode ='iterable')
    BeQu("%s%s_%s.pdb" % (output_dir, structure.name, 'clust'),
         structure, 'R', b_key ='CNT_CLUST')


nono_list = [f + '.pdb' for f in os.listdir(output_dir) if len(f) == 4]
input_list = [f for f in os.listdir(input_dir) if f.endswith('.pdb')]
if nono:
    input_list = clean_list(input_list, nono_list)


#MAIN LOOP
errors = defaultdict(list)
for (structure, file_) in izip(structure_parser(input_dir, input_list), input_list):
    print file_
    try:
        if not structure:
            errors['parsing_error'].append(file_)
            raise Exception
        if proc:
            for function in proc_funcs:
                try:
                    eval(function)(structure)
                except:
                    errors['processing_error'].append((file_, function))
                    raise Exception
        if check:
            for function in check_funcs:
                result = eval(function)(structure)
                if not result:
                    errors['structure_error'].append((file_, function))
                    raise Exception
        if write:
            try:
                output_file = output_dir + file_
                pdb_writer.set_structure(structure)
                pdb_writer.save(output_file)
            except:
                errors['pdbwrite_error'].append(file_)
                raise Exception
        if run:
            for function in run_funcs:
                try:
                    result = eval(function)(structure, output_file)
                except PDBExternalException:
                    result = [None]
                if not all(result):
                    errors['external_error'].append((file_, function))
                    raise Exception
        if derive:
            try:
                derive_externals(structure)
            except:
               _, e, tb = sys.exc_info(); print e
               errors['derive_error'].append((file_, e))
               raise Exception
        if transfer:
            try:
                transfer_coordination(structure)
            except:
               _, e, tb = sys.exc_info(); print e
               errors['transfer_error'].append((file_, e))
               raise Exception
        if dump:
            try:
                gzip_dump(structure, output_dir + file_[0:4])
            except:
               _, e, tb = sys.exc_info(); print e
               errors['pickle_error'].append((file_, e))
               raise Exception
        if color:
            try:
                color_bq(structure, bq_dir)
            except:
               _, e, tb = sys.exc_info(); print e
               errors['bq_error'].append((file_, e))
               raise Exception
    except:
        if sandbox:
            continue
        else:
            raise

#EOF
